import asyncio
from data_dyne.models.survey_model import PaperChunk
from data_dyne.rag.kb_api import KBAPI


class KBService:
    def __init__(self, semaphore=None):
        self.kb_api = KBAPI(semaphore=semaphore)

    async def get_papers(
        self,
        query: str,
        include_abstract_papers: bool = True,
        include_normal_papers: bool = True,
        merge_papers: bool = False,
        top_k = 10
    ) -> list[PaperChunk]:
        """
        根据查询获取相关论文片段

        参数:
            query (str): 搜索查询词
            include_abstract_papers (bool, optional): 是否包含摘要部分的论文片段。默认为True。
            include_normal_papers (bool, optional): 是否包含正文部分的论文片段。默认为True。
            merge_papers (bool, optional): 是否合并来自同一篇论文的片段。默认为True。

        返回:
            list[PaperChunk]: 按论文ID排序的论文片段列表。每个片段包含:
                - 论文ID
                - 标题
                - 内容
                - 相似度得分
                - 其他元数据

        说明:
            1. 当 include_abstract_papers 为True时，从论文摘要部分搜索相关片段
            2. 当 include_normal_papers 为True时，从论文正文部分搜索相关片段
            3. 将摘要和正文部分的搜索结果合并，并按论文ID排序
            4. 如果 merge_papers 为True，会将来自同一篇论文的片段合并成一个
            5. 每个片段都包含相似度得分，可用于后续重新排序
        """
        # abstract_papers = []
        # normal_papers = []
        tasks = []
        # 搜索摘要部分的论文片段 下面两种query可以放一起并发
        if include_abstract_papers:
            tasks.append(self.kb_api.search_papers_abs_with_score(query))
        else:
            tasks.append(self._fake_respone(query))
        # 搜索正文部分的论文片段
        if include_normal_papers:
            tasks.append(self.kb_api.search_papers_with_score(query))
        else:
            tasks.append(self._fake_respone(query))
        # 合并摘要和正文部分的论文片段
        results = await asyncio.gather(*tasks)

        paper_chunks = results[0] + results[1]

        # 按论文ID排序,方便后续处理
        paper_chunks.sort(key=lambda x: x.paper_id)
        if merge_papers:
            i = 0
            while i < len(paper_chunks) - 1:
                if paper_chunks[i].paper_id == paper_chunks[i + 1].paper_id:
                    paper_chunks[i].chunk_text += "\n" + paper_chunks[i + 1].chunk_text
                    paper_chunks.pop(i + 1)
                else:
                    i += 1
        # 如果top_k不为0，则只返回前top_k个
        # 去重
        paper_chunks = list({p.id: p for p in paper_chunks}.values())
        if top_k:
            paper_chunks = paper_chunks[:top_k]
        return paper_chunks
    
    async def _fake_respone(self, query: str) -> list[PaperChunk]:
        return []
    
    async def search_papers(self, query: str) -> list[PaperChunk]:
        return await self.kb_api.search_papers(query)

    async def search_papers_abs_with_score(self, query: str) -> list[PaperChunk]:
        return await self.kb_api.search_papers_abs_with_score(query)

    async def search_papers_with_score(self, query: str) -> list[PaperChunk]:
        return await self.kb_api.search_papers_with_score(query)
